import socket
import sys
import random
import time
from datetime import datetime


class ReliableUDPServer:
    def __init__(self, port, loss_rate=0.1):
        # 创建UDP套接字并绑定端口
        self.serversock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.serversock.bind(('0.0.0.0', port))
        self.loss_rate = loss_rate  # 丢包率
        self.clients = {}  # 存储客户端状态
        self.received_data = {}  # 存储接收的数据
        self.running = True

    def start(self):
        print(f"服务器启动于端口：{self.serversock.getsockname()[1]} | 等待客户端连接...")

        while self.running:
            try:
                data, addr = self.serversock.recvfrom(1024)
                self.handle_packet(data, addr)
            except KeyboardInterrupt:
                print("\n服务器关闭")
                self.running = False
                break
        self.serversock.close()

    def handle_packet(self, data, addr):
        seq_num, ack_num, flags, payload = self.parse_packet(data)

        # 处理连接请求
        if flags == 'SYN':
            self.handle_syn(seq_num, addr)
        # 处理连接确认
        elif flags == 'ACK' and addr in self.clients:
            self.handle_ack(seq_num, ack_num, addr)
        # 处理数据传输
        elif flags == 'DATA' and addr in self.clients and self.clients[addr]['connected']:
            self.handle_data(seq_num, ack_num, payload, addr)
        # 处理连接关闭
        elif flags == 'FIN' and addr in self.clients and self.clients[addr]['connected']:
            self.handle_fin(seq_num, addr)

    def handle_syn(self, seq_num, addr):
        """处理客户端SYN请求，实现三次握手"""
        print(f"收到客户端{addr}的SYN请求，seq={seq_num}")

        # 模拟三次握手 - 发送SYN+ACK
        server_isn = random.randint(1000, 5000)
        syn_ack = self.create_packet(
            seq_num=server_isn,
            ack_num=seq_num + 1,
            flags='SYN+ACK'
        )

        # 发送SYN+ACK，不实现复杂重试
        self.serversock.sendto(syn_ack, addr)
        print(f"向{addr}发送SYN+ACK，seq={server_isn}, ack={seq_num + 1}")

        # 初始化客户端状态
        self.clients[addr] = {
            'expected_seq': 1,
            'connected': False,
            'server_isn': server_isn
        }

    def handle_ack(self, seq_num, ack_num, addr):
        """处理客户端ACK，完成三次握手"""
        if addr not in self.clients:
            return

        server_isn = self.clients[addr]['server_isn']
        if ack_num == server_isn + 1:
            print(f"收到{addr}的ACK，确认连接建立，ack={ack_num}")
            self.clients[addr]['connected'] = True

            # 发送连接确认
            connected_packet = self.create_packet(
                seq_num=server_isn + 1,
                ack_num=ack_num,
                flags='CONNECTED'
            )
            self.serversock.sendto(connected_packet, addr)
            print(f"向{addr}发送连接确认，seq={server_isn + 1}")

    def handle_data(self, seq_num, ack_num, payload, addr):
        """处理客户端数据，实现GBN协议"""
        if addr not in self.clients or not self.clients[addr]['connected']:
            return

        expected_seq = self.clients[addr]['expected_seq']

        # 模拟丢包（仅对数据包）
        if random.random() < self.loss_rate:
            print(f"模拟丢包：丢弃来自{addr}的seq={seq_num}数据包")
            return

        # GBN特性：只接受期望序号的数据包
        if seq_num == expected_seq:
            print(f"收到期望包{seq_num}，来自{addr}")
            self.received_data[seq_num] = payload.decode()

            # 发送累积ACK
            ack_packet = self.create_packet(
                seq_num=0,
                ack_num=expected_seq + 1,
                flags='ACK'
            )
            self.serversock.sendto(ack_packet, addr)
            print(f"向{addr}发送ACK，期望下一个seq={expected_seq + 1}")

            # 递增期望序号
            self.clients[addr]['expected_seq'] += 1
        else:
            print(f"收到乱序包{seq_num}，期望{expected_seq}，来自{addr}")
            # 发送当前期望序号作为ACK
            ack_packet = self.create_packet(
                seq_num=0,
                ack_num=expected_seq,
                flags='ACK'
            )
            self.serversock.sendto(ack_packet, addr)

    def handle_fin(self, seq_num, addr):
        """处理客户端FIN，实现连接关闭"""
        print(f"收到{addr}的FIN，seq={seq_num}")

        # 发送ACK确认
        ack_packet = self.create_packet(
            seq_num=0,
            ack_num=seq_num + 1,
            flags='ACK'
        )
        self.serversock.sendto(ack_packet, addr)
        print(f"向{addr}发送FIN确认，ack={seq_num + 1}")

        # 发送FIN
        fin_seq = random.randint(1, 1000)
        fin_packet = self.create_packet(
            seq_num=fin_seq,
            flags='FIN'
        )
        self.serversock.sendto(fin_packet, addr)
        print(f"向{addr}发送FIN，seq={random.randint(1, 1000)}")

        # 设置超时等待客户端的ACK
        try:
            self.serversock.settimeout(3.0)
            data, addr = self.serversock.recvfrom(1024)

            s_seq, s_ack, flags, s_data = self.parse_packet(data)
            if flags == 'ACK' and s_ack == fin_seq + 1:
                print(f"[SERVER] 收到{addr}的最终ACK，ack={s_ack} (第四次挥手)")
        except socket.timeout:
            print(f"[SERVER] 等待客户端ACK超时")

        # 标记客户端为关闭
        if addr in self.clients:
            self.clients[addr]['connected'] = False
            print(f"客户端{addr}连接已关闭")

    def create_packet(self, seq_num, ack_num=0, flags='', data=''):
        """创建数据包，格式：seq|ack|flags|data"""
        return f"{seq_num}|{ack_num}|{flags}|{data}".encode()

    def parse_packet(self, packet):
        try:
            parts = packet.decode().split('|', 3)
            seq_num = int(parts[0])
            ack_num = int(parts[1])
            flags = parts[2]
            data = parts[3].encode() if len(parts) > 3 else b''  # 保持为字节类型
            return seq_num, ack_num, flags, data
        except:
            return 0, 0, '', b''


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("使用: python server.py <port>")
        sys.exit(1)

    port = int(sys.argv[1])
    server = ReliableUDPServer(port, loss_rate=0.1)
    # server = ReliableUDPServer(3083, loss_rate=0.2)
    server.start()